#!/usr/bin/env python3

import rospy
import numpy as np
from sensor_msgs.msg import PointCloud2, PointField
from sensor_msgs import point_cloud2
from visualization_msgs.msg import Marker, MarkerArray
from sklearn.cluster import DBSCAN
import open3d as o3d
import struct

class LivoxClusteringNode:
    def __init__(self):
        rospy.init_node('livox_clustering_node', anonymous=True)

        # 订阅 Livox 点云
        self.sub = rospy.Subscriber('/livox/lidar', PointCloud2, self.pointcloud_callback)

        # 发布聚类后的点云（用于 RViz 显示不同颜色）
        self.clustered_pc_pub = rospy.Publisher('/clustered_pointcloud', PointCloud2, queue_size=1)

        # 发布聚类边界框（可选：显示包围盒）
        self.marker_pub = rospy.Publisher('/cluster_bounding_boxes', MarkerArray, queue_size=10)

        self.rate = rospy.Rate(10)  # 10 Hz 处理频率

        rospy.loginfo("Livox Clustering Node Started.")

    def pointcloud2_to_xyz_array(self, pc2_msg):
        """将 PointCloud2 消息转为 Nx3 的 numpy 数组 (x, y, z)"""
        gen = point_cloud2.read_points(pc2_msg, field_names=("x", "y", "z"), skip_nans=True)
        points = np.array(list(gen), dtype=np.float32)
        return points

    def cluster_points(self, points):
        """使用 DBSCAN 聚类，并根据穿越机点数过滤"""
        if len(points) == 0:
            return np.array([]), 0

        # 根据穿越机尺寸调整参数
        db = DBSCAN(eps=0.28, min_samples=5).fit(points)  # 更小的 eps 和更少的 min_samples
        labels = db.labels_
        n_clusters = len(set(labels)) - (1 if -1 in labels else 0)

        # 过滤掉点数太多或太少的聚类（例如只保留 5~30 个点）
        filtered_labels = np.copy(labels)
        unique_labels = set(labels)
        if -1 in unique_labels:
            unique_labels.remove(-1)

        for label in unique_labels:
            cluster_points = points[labels == label]
            if not (5 <= len(cluster_points) <= 35):  # 只保留符合穿越机点数范围的聚类
                filtered_labels[labels == label] = -1  # 设为噪声

        return filtered_labels, n_clusters

    def create_colored_pointcloud(self, points, labels):
        """为每个聚类生成带颜色的 PointCloud2 消息"""
        fields = [
            PointField('x', 0, PointField.FLOAT32, 1),
            PointField('y', 4, PointField.FLOAT32, 1),
            PointField('z', 8, PointField.FLOAT32, 1),
            PointField('rgb', 12, PointField.UINT32, 1)
        ]

        # 颜色表（最多10种）
        colors = [
            [255, 0, 0],    # 红
            [0, 255, 0],    # 绿
            [0, 0, 255],    # 蓝
            [255, 255, 0],
            [255, 0, 255],
            [0, 255, 255],
            [128, 0, 0],
            [0, 128, 0],
            [0, 0, 128],
            [128, 128, 0]
        ]

        # 构造二进制数据（bytes）
        cloud_data = []
        for i, point in enumerate(points):
            x, y, z = point[0], point[1], point[2]
            label = labels[i]

            if label == -1:
                r, g, b = 0, 0, 0  # 黑色表示噪声
            else:
                c = colors[label % len(colors)]
                r, g, b = int(c[0]), int(c[1]), int(c[2])

            # 将 RGB 打包为 UINT32
            rgb = (r << 16) | (g << 8) | b
            # 注意：这里不要转换成 float，而是直接用 I 格式打包成 4 字节
            packed_data = struct.pack('fffI', x, y, z, rgb)  # 4个字段：3×float + 1×uint32
            cloud_data.append(packed_data)

        header = rospy.Header()
        header.stamp = rospy.Time.now()
        header.frame_id = "livox_frame"

        pc2_msg = PointCloud2()
        pc2_msg.header = header
        pc2_msg.height = 1
        pc2_msg.width = len(points)
        pc2_msg.is_dense = True  # 假设无无效点
        pc2_msg.is_bigendian = False
        pc2_msg.fields = fields
        pc2_msg.point_step = 16  # 4*4=16 字节每点
        pc2_msg.row_step = pc2_msg.point_step * pc2_msg.width
        pc2_msg.data = b''.join(cloud_data)  # 合并所有点的二进制数据

        return pc2_msg

    def publish_bounding_boxes(self, points, labels):
        """发布每个聚类的包围盒 Marker"""
        marker_array = MarkerArray()
        unique_labels = set(labels)
        if -1 in unique_labels:
            unique_labels.remove(-1)

        # 预定义颜色表 (r, g, b)，避免使用 plt
        colors = [
            [1.0, 0.0, 0.0],  # 红
            [0.0, 1.0, 0.0],  # 绿
            [0.0, 0.0, 1.0],  # 蓝
            [1.0, 1.0, 0.0],  # 黄
            [1.0, 0.0, 1.0],  # 品红
            [0.0, 1.0, 1.0],  # 青
            [1.0, 0.5, 0.0],  # 橙
            [0.5, 0.0, 1.0],  # 紫
            [0.0, 1.0, 0.5],
            [0.8, 0.4, 0.2]
        ]

        for i, label in enumerate(unique_labels):
            cluster_pts = points[labels == label]
            if len(cluster_pts) < 5:
                continue

            min_pt = np.min(cluster_pts, axis=0)
            max_pt = np.max(cluster_pts, axis=0)
            center = (min_pt + max_pt) / 2.0
            size = max_pt - min_pt

            marker = Marker()
            marker.header.frame_id = "livox_frame"
            marker.header.stamp = rospy.Time.now()
            marker.ns = "clusters"
            marker.id = i
            marker.type = Marker.CUBE
            marker.action = Marker.ADD
            marker.pose.position.x = center[0]
            marker.pose.position.y = center[1]
            marker.pose.position.z = center[2]
            marker.scale.x = size[0]
            marker.scale.y = size[1]
            marker.scale.z = size[2]

            # 使用预定义颜色循环
            color = colors[i % len(colors)]
            marker.color.r = color[0]
            marker.color.g = color[1]
            marker.color.b = color[2]
            marker.color.a = 0.4  # 半透明

            marker_array.markers.append(marker)

        self.marker_pub.publish(marker_array)

    def pointcloud_callback(self, msg):
        """点云回调函数"""
        try:
            points = self.pointcloud2_to_xyz_array(msg)
            if len(points) == 0:
                return
            
            # === 新增：只保留 Z 轴在 [-1, 1] 范围内的点 ===
            z_min, z_max = -0.5, 0.5  # 可根据实际安装高度调整
            valid_indices = np.where((points[:, 2] >= z_min) & (points[:, 2] <= z_max))[0]
            points = points[valid_indices]

            # === 新增：只保留距离小于等于 7 米的点 ===
            max_distance = 5.0  # 单位为米
            distances = np.linalg.norm(points[:, :2], axis=1)  # 计算 XY 平面距离（也可用三维距离 np.linalg.norm(points, axis=1)）
            valid_indices = np.where(distances <= max_distance)[0]
            points = points[valid_indices]

            # 聚类
            labels, n_clusters = self.cluster_points(points)
            rospy.loginfo_throttle(5, f"Detected {n_clusters} clusters")

            # 发布着色点云
            colored_pc = self.create_colored_pointcloud(points, labels)
            self.clustered_pc_pub.publish(colored_pc)

            # 发布包围盒（可选）
            self.publish_bounding_boxes(points, labels)

        except Exception as e:
            rospy.logerr(f"Error in processing point cloud: {e}")

    def run(self):
        rospy.spin()


if __name__ == '__main__':
    try:
        node = LivoxClusteringNode()
        node.run()
    except rospy.ROSInterruptException:
        pass